# coding=utf-8
from numpy import *
 
def loadDataSet():
    """创建一个用于测试的简单的数据集"""
    return [[1, 3, 4], [2, 3, 5], [1, 2, 3, 5], [2, 5]]

def createC1(dataSet):
    """
	构建初始候选项集的列表，即所有候选项集只包含一个元素，
	C1是大小为1的所有候选项集的集合
    """
    C1 = []
    for transaction in dataSet:
        for item in transaction:
            if not [item] in C1:
                C1.append([item])
    C1.sort()
    return map(frozenset, C1)

def scanD(D, Ck, minSupport):
    """
	计算Ck中的项集在数据集D(记录或者transactions)中的支持度，
	返回满足最小支持度的项集的集合，和所有项集支持度信息的字典。
    """
    ssCnt = {}
    for tid in D:
        for can in Ck:
            if can.issubset(tid):
                ssCnt[can] = ssCnt.get(can, 0) + 1
    numItems = float(len(D))
    retList = []
    supportData = {}
    for key in ssCnt:
        support = ssCnt[key] / numItems
        if support >= minSupport:
            retList.insert(0, key)
        supportData[key] = support
    return retList, supportData

def aprioriGen(Lk, k):
    """
	由初始选项集的集合Lk生成新的生成候选项集，
	k表示生成的新项集中所含有的元素的个数
    """
    retList = []
    lenLk = len(Lk)
    for i in range(lenLk):
        for j in range(i + 1, lenLk):
            # 前k-2项相同时，将两个集合合并
            L1 = list(Lk[i])[:k-2]; L2 = list(Lk[j])[:k-2]
            L1.sort(); L2.sort()
            if L1 == L2:
                retList.append(Lk[i] | Lk[j])
    return retList

def apriori(dataSet, minSupport=0.5):
    # 构建初始候选项集C1
    C1 = createC1(dataSet)
    # 将dataset集合化，以满足scan的格式要求
    D = map(set, dataSet)
    # 构建初始的频繁项集，即所有项集只有一个元素
    L1, supportData = scanD(D, C1, minSupport)
    L = [L1]
    # 最初的L1中的每个项集含有一个元素，新生成的
    # 项集应该含有2个元素，所以k=2
    k = 2
    while (len(L[k-2]) > 0):
        Ck = aprioriGen(L[k-2], k)
        Lk, supK = scanD(D, Ck, minSupport)
        supportData.update(supK)
        L.append(Lk)
        k += 1
    return L, supportData

def generateRules(L, supportData, minConf=0.7):
    """
	关联规则生成
    """
    bigRuleList = []
    for i in range(1, len(L)):
        for freqSet in L[i]:
            H1 = [frozenset([item]) for item in freqSet]
            if (i > 1):
                # 三个及以上元素的集合
                rulesFromConseq(freqSet, H1, supportData, bigRuleList, minConf)
            else:
                # 两个元素的集合
                calcConf(freqSet, H1, supportData, bigRuleList, minConf)
    return bigRuleList

def calcConf(freqSet, H, supportData, brl, minConf=0.7):
    ''' 对候选规则集进行评估 '''
    prunedH = []
    for conseq in H:
        conf = supportData[freqSet] / supportData[freqSet - conseq]
        if conf >= minConf:
            print freqSet - conseq, '-->', conseq, 'conf:', conf
            brl.append((freqSet - conseq, conseq, conf))
            prunedH.append(conseq)
    return prunedH

def rulesFromConseq(freqSet, H, supportData, brl, minConf=0.7):
    ''' 生成候选规则集 '''
    m = len(H[0])
    if (len(freqSet) > (m + 1)):
        Hmpl = aprioriGen(H, m + 1)
        Hmpl = calcConf(freqSet, Hmpl, supportData, brl, minConf)
        if (len(Hmpl) > 1):
            rulesFromConseq(freqSet, Hmpl, supportData, brl, minConf)

def generateRules2(L, supportData, minConf=0.7):
    """
	在i>1时，将对H1调用calcConf()的过程加上
    """
    bigRuleList = []
    for i in range(1, len(L)):
        for freqSet in L[i]:
            H1 = [frozenset([item]) for item in freqSet]
            if (i > 1):
                # 三个及以上元素的集合
                H1 = calcConf(freqSet, H1, supportData, bigRuleList, minConf)
                rulesFromConseq(freqSet, H1, supportData, bigRuleList, minConf)
            else:
                # 两个元素的集合
                calcConf(freqSet, H1, supportData, bigRuleList, minConf)
    return bigRuleList

def generateRules3(L, supportData, minConf=0.7):
    """
	当i=1时的else部分并没有独特的逻辑，这个if语句可以合并，
	然后再修改rulesFromConseq()函数，保证其会调用calcConf(freqSet, H1, …)
    """
    bigRuleList = []
    for i in range(1, len(L)):
        for freqSet in L[i]:
            H1 = [frozenset([item]) for item in freqSet]
            rulesFromConseq2(freqSet, H1, supportData, bigRuleList, minConf)
    return bigRuleList
 
def rulesFromConseq2(freqSet, H, supportData, brl, minConf=0.7):
    m = len(H[0])
    if (len(freqSet) > m): # 判断长度改为 > m，这时即可以求H的可信度
        Hmpl = calcConf(freqSet, H, supportData, brl, minConf)
        if (len(Hmpl) > 1): # 判断求完可信度后是否还有可信度大于阈值的项用来生成下一层H
            Hmpl = aprioriGen(Hmpl, m + 1)
            rulesFromConseq2(freqSet, Hmpl, supportData, brl, minConf) # 递归计算，不变

def rulesFromConseq3(freqSet, H, supportData, brl, minConf=0.7):
    """
	消除rulesFromConseq2()函数中的递归项。这个递归纯粹是偷懒的结果，
	没有简化任何逻辑和增加任何可读性，可以直接用一个循环代替
    """
    m = len(H[0])
    while (len(freqSet) > m): # 判断长度 > m，这时即可求H的可信度
        H = calcConf(freqSet, H, supportData, brl, minConf)
        if (len(H) > 1): # 判断求完可信度后是否还有可信度大于阈值的项用来生成下一层H
            H = aprioriGen(H, m + 1)
            m += 1
        else: # 不能继续生成下一层候选关联规则，提前退出循环
            break

dataSet = loadDataSet()
print dataSet
print '-'*20
C1 = createC1(dataSet)
D = map(set, dataSet)
print D
print '-'*20
L1, suppDat = scanD(D, C1, 0.5)
print L1
print '-'*20
L, suppData = apriori(dataSet)
print L
print '-'*20
# L, suppData = apriori(dataSet, minSupport=0.7)
# print L
rules = generateRules(L, suppData, minConf=0.7)
print rules
print '-'*20
rules = generateRules(L, suppData, minConf=0.5)
print rules
print '-'*20